mesa_llm/llm_agent.py (3)

105-113: ⚠️ Potential issue | 🟠 Major

Preserve all tool-call results in memory.

The dict comprehension collapses every tool result into the same name / response keys, so only the last tool call survives in memory when a plan executes more than one tool. Store a list under a stable key instead of flattening the response list.

🛠️ Suggested fix

         await self.memory.aadd_to_memory(
             type="action",
             content={
-                k: v
-                for tool_call in tool_call_resp
-                for k, v in tool_call.items()
-                if k not in ["tool_call_id", "role"]
+                "tool_calls": [
+                    {
+                        k: v
+                        for k, v in tool_call.items()
+                        if k not in {"tool_call_id", "role"}
+                    }
+                    for tool_call in tool_call_resp
+                ]
             },
         )
...
         self.memory.add_to_memory(
             type="action",
             content={
-                k: v
-                for tool_call in tool_call_resp
-                for k, v in tool_call.items()
-                if k not in ["tool_call_id", "role"]
+                "tool_calls": [
+                    {
+                        k: v
+                        for k, v in tool_call.items()
+                        if k not in {"tool_call_id", "role"}
+                    }
+                    for tool_call in tool_call_resp
+                ]
             },
         )

Also applies to: 133-141

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@mesa_llm/llm_agent.py` around lines 105 - 113, The current
memory.aadd_to_memory call flattens multiple tool_call_resp dicts into a single
dict, losing earlier tool results; instead, change the payload to store the
entire list under a stable key (e.g., "tool_call_results" or "tool_calls") so
all entries are preserved (reference: memory.aadd_to_memory and variable
tool_call_resp), and make the same change to the similar block handling
tool_call_resp at the later section (lines 133-141) so both places append the
full list rather than collapsing into one dict.

---

344-349: ⚠️ Potential issue | 🔴 Critical

Balance pre/post-step hooks exactly once, even on errors.

astep() opens an async pre/post pair and then calls self.step(), but overridden step() methods are wrapped below with their own pre/post hooks. For subclasses that only implement step, that double-opens the lifecycle and leaves memory bookkeeping unbalanced. The wrappers also skip the post hook when user code raises.

🛠️ Suggested fix

     async def astep(self) -> None:
-        await self.apre_step()
-
         if hasattr(self, "step") and self.__class__.step != LLMAgent.step:
             self.step()
+            return
 
-        await self.apost_step()
+        await self.apre_step()
+        try:
+            return
+        finally:
+            await self.apost_step()
...
             def wrapped(self, *args: Any, **kwargs: Any) -> Any:
-                LLMAgent.pre_step(self, *args, **kwargs)
-                result = user_step(self, *args, **kwargs)
-                LLMAgent.post_step(self, *args, **kwargs)
-                return result
+                LLMAgent.pre_step(self)
+                try:
+                    return user_step(self, *args, **kwargs)
+                finally:
+                    LLMAgent.post_step(self)
...
             async def awrapped(self, *args: Any, **kwargs: Any) -> Any:
                 await self.apre_step()
-                result = await user_astep(self, *args, **kwargs)
-                await self.apost_step()
-                return result
+                try:
+                    return await user_astep(self, *args, **kwargs)
+                finally:
+                    await self.apost_step()

Also applies to: 360-378

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@mesa_llm/llm_agent.py` around lines 344 - 349, The pre/post-step lifecycle is
being opened twice for subclasses that override step and may skip post on
exceptions; fix by calling apre_step() once, executing self.step() inside a
try/finally, and always calling apost_step() in the finally so the pair is
balanced even on errors; update the same pattern for the other wrapper region
(the methods around lines 360-378) and ensure any wrapper implementations of
step do not themselves call apre_step/apost_step so the lifecycle is only
managed in the outer caller.

---

229-249: ⚠️ Potential issue | 🟠 Major

Replace self.model.step with self.model.steps in two locations.

Model.step() is the simulation method; Model.steps is the integer counter tracking executed steps. Using self.model.step assigns a bound method object to step: int, violating the Observation type contract and causing runtime errors.

Affected locations

• Line 239: step: int = self.model.step → step: int = self.model.steps
• Line 253: step: int = self.model.step → step: int = self.model.steps

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@mesa_llm/llm_agent.py` around lines 229 - 249, The code assigns a bound
method object to the Observation.step field by using self.model.step instead of
the integer counter; update both occurrences to use the integer counter property
self.model.steps (e.g., in agenerate_obs and the corresponding generate_obs
method) so step: int = self.model.steps and the returned Observation(step=step,
...) uses an int, not a bound method.